import json
import jieba
from collections import defaultdict
import re
from datetime import datetime
import os

class InvertedIndex:
    def __init__(self):
        self.inverted_index = defaultdict(list)
        self.documents = {}
        
    def preprocess_text(self, text):
        if not isinstance(text, str):
            text = str(text)
        # 移除标点符号和特殊字符
        text = re.sub(r'[^\w\s]', '', text)
        # 使用jieba进行分词
        words = jieba.cut(text)
        return [word for word in words if word.strip()]
    
    def load_existing_index(self):
        """尝试加载最新的索引文件"""
        try:
            # 检查index_files目录是否存在
            if not os.path.exists('index_files'):
                return False
                
            # 获取最新的索引文件
            index_files = [f for f in os.listdir('index_files') if f.startswith('inverted_index_')]
            if not index_files:
                return False
                
            latest_index = max(index_files)
            latest_docs = latest_index.replace('inverted_index_', 'documents_')
            
            # 加载索引
            with open(os.path.join('index_files', latest_index), 'r', encoding='utf-8') as f:
                index_data = json.load(f)
                self.inverted_index = defaultdict(list, {
                    word: data['document_ids']
                    for word, data in index_data['index'].items()
                })
                
            # 加载文档
            with open(os.path.join('index_files', latest_docs), 'r', encoding='utf-8') as f:
                docs_data = json.load(f)
                self.documents = docs_data['documents']
                
            print(f"\n已加载现有索引文件: {latest_index}")
            print(f"索引包含 {len(self.inverted_index)} 个词条")
            print(f"文档数量: {len(self.documents)} 个")
            return True
            
        except Exception as e:
            print(f"加载索引文件时出错: {e}")
            return False
    
    def build_index(self, json_file, force_rebuild=False):
        """
        建立倒排索引
        :param json_file: JSON文件路径
        :param force_rebuild: 是否强制重建索引，即使已存在索引文件
        """
        try:
            # 如果不是强制重建且成功加载了现有索引，则直接返回
            if not force_rebuild and self.load_existing_index():
                return
                
            print("开始建立新的倒排索引...")
            
            # 读取JSON文件
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 处理字典形式的数据
            if not isinstance(data, dict):
                raise ValueError("JSON文件格式不正确，应为字典对象")
            
            # 清空现有索引（如果有的话）
            self.inverted_index.clear()
            self.documents.clear()
            
            # 为每个文档建立索引
            for doc_id, item in data.items():
                # 存储原始文档
                self.documents[doc_id] = item
                
                # 获取文本内容（同时索引标题和正文）
                text = item.get('title', '') + ' ' + item.get('text', '')
                
                # 对文本进行预处理和分词
                words = self.preprocess_text(text)
                
                # 建立倒排索引
                for word in words:
                    if doc_id not in self.inverted_index[word]:
                        self.inverted_index[word].append(doc_id)
                        
            print(f"成功索引了 {len(self.documents)} 个文档")
            print(f"词汇表大小: {len(self.inverted_index)} 个词")
            
            # 保存倒排索引和文档信息
            self.save_index()
            
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
            raise
        except Exception as e:
            print(f"处理文件时出错: {e}")
            raise

    def save_index(self):
        """保存索引到文件"""
        # 创建输出目录
        if not os.path.exists('index_files'):
            os.makedirs('index_files')
        
        # 获取当前时间戳
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 保存倒排索引
        inverted_index_file = f'index_files/inverted_index_{timestamp}.json'
        with open(inverted_index_file, 'w', encoding='utf-8') as f:
            # 将defaultdict转换为普通dict
            index_dict = {
                'word_count': len(self.inverted_index),
                'document_count': len(self.documents),
                'index': {
                    word: {
                        'document_count': len(doc_ids),
                        'document_ids': doc_ids
                    }
                    for word, doc_ids in sorted(self.inverted_index.items())
                }
            }
            json.dump(index_dict, f, ensure_ascii=False, indent=2)
        print(f"\n倒排索引已保存至: {inverted_index_file}")
        
        # 保存文档信息
        docs_file = f'index_files/documents_{timestamp}.json'
        with open(docs_file, 'w', encoding='utf-8') as f:
            json.dump({
                'document_count': len(self.documents),
                'documents': self.documents
            }, f, ensure_ascii=False, indent=2)
        print(f"文档信息已保存至: {docs_file}")
        
        # 生成人类可读的倒排索引报告
        report_file = f'index_files/index_report_{timestamp}.txt'
        with open(report_file, 'w', encoding='utf-8') as f:
            f.write("倒排索引报告\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"总文档数: {len(self.documents)}\n")
            f.write(f"总词汇量: {len(self.inverted_index)}\n\n")
            f.write("词汇表及其出现的文档：\n")
            f.write("-" * 50 + "\n\n")
            
            # 按文档频率降序排序词汇
            sorted_words = sorted(self.inverted_index.items(), 
                                key=lambda x: len(x[1]), 
                                reverse=True)
            
            for word, doc_ids in sorted_words:
                doc_count = len(doc_ids)
                f.write(f"词语: {word}\n")
                f.write(f"出现在 {doc_count} 个文档中: {', '.join(doc_ids)}\n")
                # 显示每个文档的标题
                f.write("文档标题:\n")
                for doc_id in doc_ids:
                    title = self.documents[doc_id].get('title', 'N/A')
                    f.write(f"  - {doc_id}: {title}\n")
                f.write("\n")
                
        print(f"索引报告已保存至: {report_file}")
    
    def search(self, query):
        # 对查询进行分词
        query_words = self.preprocess_text(query)
        
        if not query_words:
            return []
        
        # 获取包含查询词的文档
        matching_docs = set()
        first_word = True
        
        for word in query_words:
            if word in self.inverted_index:
                if first_word:
                    matching_docs = set(self.inverted_index[word])
                    first_word = False
                else:
                    # 对多个查询词取交集
                    matching_docs &= set(self.inverted_index[word])
        
        # 返回匹配的原始文档
        results = []
        for doc_id in matching_docs:
            doc = self.documents[doc_id]
            results.append({
                'id': doc_id,
                'title': doc.get('title', 'N/A'),
                'url': doc.get('url', 'N/A'),
                'text': doc.get('text', 'N/A'),
                'pubtime': doc.get('pubtime', 'N/A')
            })
        
        # 创建output_search目录（如果不存在）
        if not os.path.exists('output_search'):
            os.makedirs('output_search')
        
        # 生成搜索结果文件名（使用查询关键词）
        filename = '_'.join(query_words)  # 将多个关键词用下划线连接
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        output_file = f'output_search/{filename}_{timestamp}.json'
        
        # 保存搜索结果为JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump({
                'query': query,
                'timestamp': timestamp,
                'result_count': len(results),
                'results': results
            }, f, ensure_ascii=False, indent=2)
            
        print(f"\n搜索结果已保存至: {output_file}")
        
        return results

    def _tokenize_boolean_query(self, query):
        """将布尔查询分词，保留操作符和括号"""
        # 替换操作符为标准格式
        query = query.replace('AND', 'and').replace('OR', 'or').replace('NOT', 'not')
        # 在括号周围添加空格
        query = query.replace('(', ' ( ').replace(')', ' ) ')
        # 分割成token
        tokens = query.split()
        return tokens

    def _evaluate_boolean_expression(self, tokens):
        """计算布尔表达式"""
        def get_precedence(op):
            if op == 'not':
                return 3
            if op == 'and':
                return 2
            if op == 'or':
                return 1
            return 0

        def apply_operator(operators, operands):
            if not operators:
                return
            op = operators.pop()
            if op == 'not':
                operand = operands.pop()
                # 获取所有文档的集合
                all_docs = set(doc_id for doc_ids in self.inverted_index.values() for doc_id in doc_ids)
                # NOT操作：所有文档减去包含该词的文档
                operands.append(all_docs - operand)
            else:
                right = operands.pop()
                left = operands.pop()
                if op == 'and':
                    operands.append(left & right)
                elif op == 'or':
                    operands.append(left | right)

        operators = []
        operands = []
        i = 0
        
        while i < len(tokens):
            token = tokens[i]
            
            if token == '(':
                operators.append(token)
            elif token == ')':
                while operators and operators[-1] != '(':
                    apply_operator(operators, operands)
                operators.pop()  # 移除 '('
            elif token in ('and', 'or', 'not'):
                while (operators and operators[-1] != '(' and 
                       get_precedence(operators[-1]) >= get_precedence(token)):
                    apply_operator(operators, operands)
                operators.append(token)
            else:
                # 对普通词进行分词并获取文档集合
                words = self.preprocess_text(token)
                doc_set = set()
                for word in words:
                    if word in self.inverted_index:
                        if not doc_set:
                            doc_set = set(self.inverted_index[word])
                        else:
                            # 对同一个词组中的词取交集
                            doc_set &= set(self.inverted_index[word])
                operands.append(doc_set)
            i += 1

        while operators:
            apply_operator(operators, operands)

        return operands[0] if operands else set()

    def boolean_search(self, query):
        """执行布尔搜索"""
        # 分词得到tokens
        tokens = self._tokenize_boolean_query(query)
        
        # 计算布尔表达式
        matching_docs = self._evaluate_boolean_expression(tokens)
        
        # 返回匹配的原始文档
        results = []
        for doc_id in matching_docs:
            doc = self.documents[doc_id]
            results.append({
                'id': doc_id,
                'title': doc.get('title', 'N/A'),
                'url': doc.get('url', 'N/A'),
                'text': doc.get('text', 'N/A'),
                'pubtime': doc.get('pubtime', 'N/A')
            })
        
        # 创建output_search目录（如果不存在）
        if not os.path.exists('output_search'):
            os.makedirs('output_search')
        
        # 生成搜索结果文件名
        safe_query = re.sub(r'[^\w\s]', '_', query)  # 将特殊字符替换为下划线
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        output_file = f'output_search/boolean_search_{safe_query}_{timestamp}.json'
        
        # 保存搜索结果为JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump({
                'query': query,
                'timestamp': timestamp,
                'result_count': len(results),
                'results': results
            }, f, ensure_ascii=False, indent=2)
            
        print(f"\n搜索结果已保存至: {output_file}")
        
        return results

def main():
    # 创建倒排索引实例
    index = InvertedIndex()
    
    try:
        # 建立索引
        print("正在处理索引...")
        # 默认不强制重建索引
        index.build_index('url_list.json', force_rebuild=False)
        print("索引处理完成！")
        
        # 交互式搜索
        while True:
            print("\n请选择搜索模式：")
            print("1. 普通搜索")
            print("2. 布尔搜索（支持 AND, OR, NOT, 括号）")
            print("q. 退出")
            
            choice = input("请输入选择: ")
            
            if choice.lower() == 'q':
                break
            
            query = input("\n请输入搜索关键词: ")
            if query.lower() == 'q':
                break
            
            if choice == '1':
                results = index.search(query)
            elif choice == '2':
                results = index.boolean_search(query)
            else:
                print("无效的选择，请重试")
                continue
            
            print(f"\n找到 {len(results)} 条匹配结果：")
            
            for i, doc in enumerate(results, 1):
                print(f"\n结果 {i}:")
                print(f"ID: {doc['id']}")
                print(f"发布时间: {doc['pubtime']}")
                print(f"标题: {doc['title']}")
                print(f"URL: {doc['url']}")
                print("文本预览: " + doc['text'][:200] + ("..." if len(doc['text']) > 200 else ""))
    
    except KeyboardInterrupt:
        print("\n程序已终止")
    except Exception as e:
        print(f"发生错误: {e}")

if __name__ == "__main__":
    main() 